from typing import Dict,Any,List
from multiprocessing import Queue,Process,Pool,Manager
from queue import Empty
pool:List[Process]=[]
MAX_PROCESS=4
MODULE_NAME="normal"
print(f"Module {MODULE_NAME} loads.")
manager=Manager()
current_task=manager.dict()
task_data=manager.dict()
pending=manager.Queue()
result=manager.Queue()
import threading,signal
should_stop=threading.Event()
def add_task(rid):
    from core.judgers.judgers_main import get_data,set_fail
    global pending,task_data
    if should_stop.is_set():
        set_fail(rid)
        return
    task_data[rid]=get_data(rid)
    pending.put(rid)

submitter=None
initialized=False

def submitter_mainloop():
    from core.judgers.judgers_main import submit_result
    global result,should_stop
    print("Submitter launched.")
    while not should_stop.is_set():
        try:
            res=result.get(timeout=1.5)
        except Empty:
            #print("No submission found, continuing.")
            #print(should_stop.is_set()," SET?")
            continue
        except:
            print(f"Submitter of {MODULE_NAME} quits.")
            break
        submit_result(res)
    print(f"Submitter {MODULE_NAME} quits")

def at_exit():
    from core.judgers.judgers_main import set_fail,submit_result
    print(f"Moudule {MODULE_NAME}: scheduled exit.")
    global pool,current_task,pending,result,should_stop
    if should_stop.is_set():
        print("Clear procedure executed.")
        return
    should_stop.set()
    for i in pool:
        if i.is_alive()==False:
            i.close()
            continue
        i.kill()
        i.close()
    print("Ss ",should_stop.is_set())
    print(f"{MAX_PROCESS} worker(s) terminated")
    for i in current_task:
        set_fail[current_task[i]]
    
    while True:
        try:
            rid=pending.get_nowait()
        except:
            break
        set_fail(rid)
    
    while True:
        try:
            res=result.get_nowait()
        except:
            break
        submit_result(res)
    
    print(f"{MODULE_NAME} ended successfully")

from .worker import worker
def judger_init():
    global initialized
    if initialized==True:
        print(f"Initialized module {MODULE_NAME}")
        return
    initialized=True
    should_stop.clear()
    global pool,submitter
    for i in range(MAX_PROCESS):
        print(f"Created worker #{i}")
        q=Process(target=worker,args=(i,current_task,pending,result,task_data))
        pool.append(q)
    print(f"Launch {MAX_PROCESS} process...")
    for i in pool:
        i.start()
    print(f"Done.")

    print("Start submitter")
    import threading
    submitter=threading.Thread(target=submitter_mainloop)
    submitter.start()
    print(f"Completed {MODULE_NAME}")


    
